{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2851a63e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "import cv2\n",
    "import json\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f5f0434",
   "metadata": {},
   "outputs": [],
   "source": [
    "def im_show(image):\n",
    "    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "    plt.imshow(image)\n",
    "    plt.show()\n",
    "\n",
    "def read_point_in(input_path):\n",
    "    frame_2_points = {}\n",
    "    with open(input_path, 'r') as fin:\n",
    "        for line in fin.readlines():\n",
    "            items = line.strip().split(',')\n",
    "            frame_idx = int(items[0])\n",
    "            point_info = items[1:]\n",
    "            if frame_idx in frame_2_points:\n",
    "                frame_2_points[frame_idx].append(point_info)\n",
    "            else:\n",
    "                frame_2_points[frame_idx] = [point_info]\n",
    "    return frame_2_points\n",
    "\n",
    "def write_json_out(output_path, out_frame_2_person):\n",
    "    frame_2_person = dict(sorted(out_frame_2_person.items()))\n",
    "    with open(output_path, 'w') as fout:\n",
    "        for frame_idx, persons in frame_2_person.items():\n",
    "            out_dict = {\"frame_id\":frame_idx, \"frame_pts\":0}\n",
    "            coor_list = []\n",
    "            for person in persons:\n",
    "                num_txt, l, t, r, b = person\n",
    "                num, team = num_txt.split('_')\n",
    "                num = int(num)\n",
    "                coor = {'coord':[l, t, r, b], \"name\":num_txt, \"number\":num, \"team\":team}\n",
    "                coor_list.append(coor)\n",
    "            out_dict['data'] = coor_list\n",
    "            line = json.dumps(out_dict)\n",
    "            fout.write(\"{}\\n\".format(line))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "211456be",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "input_video = '/mnt/nas/03_自由视角原始素材/202303_ZCAM相机_CBA-东莞_23-24/广东vs浙江_fv2_test/view_014_mask_ground.mp4'\n",
    "output_video = './fv_out.mp4'\n",
    "\n",
    "\n",
    "reader = cv2.VideoCapture(input_video)\n",
    "width = int(reader.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "height = int(reader.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "frame_id = -1\n",
    "more = True\n",
    "writer = cv2.VideoWriter(output_video, cv2.VideoWriter_fourcc(*\"mp4v\"), 25, (width, height))\n",
    "\n",
    "frame_v0_2_points = read_point_in(\"view_out/000_to_014_yolox_x.txt\")\n",
    "frame_v25_2_points = read_point_in(\"view_out/025_to_014_yolox_x.txt\")\n",
    "frame_v35_2_points = read_point_in(\"view_out/035_to_014_yolox_x.txt\")\n",
    "\n",
    "out_frame_id_2_person = {}\n",
    "\n",
    "frame_2_points_list = [frame_v0_2_points, frame_v25_2_points, frame_v35_2_points]\n",
    "while more:\n",
    "    more, frame = reader.read()\n",
    "    frame_id += 1\n",
    "#     if(frame_id >= 5):\n",
    "#         break\n",
    "        \n",
    "    if(frame_id % 20 == 0):\n",
    "        print(\"== handle \", frame_id)\n",
    "        \n",
    "    num_txt_2_person_list = {}\n",
    "    num_person_list = []\n",
    "    \n",
    "    for frame_2_points in frame_2_points_list:\n",
    "        if frame_id in frame_2_points:\n",
    "            persons = frame_2_points[frame_id]\n",
    "        else:\n",
    "            person = []\n",
    "\n",
    "        \n",
    "        for person in persons:\n",
    "            #print(person)\n",
    "            num_txt, r, b, w, h = person\n",
    "            r, b, w, h = int(r), int(b), int(w), int(h)\n",
    "            t = b-h\n",
    "            l = r-w\n",
    "            \n",
    "            if(num_txt in num_txt_2_person_list):\n",
    "                num_txt_2_person_list[num_txt].append((t, l, b, r))\n",
    "            else:\n",
    "                num_txt_2_person_list[num_txt] = [(t, l, b, r)]\n",
    "            cv2.circle(frame, (r,b), 5, (255, 255, 0), 5)\n",
    "            #cv2.putText(frame, num_txt, (r, b), cv2.FONT_HERSHEY_COMPLEX, 3, (255, 0, 0), 3)\n",
    "            \n",
    "            #cv2.rectangle(frame, (l, t), (r, b), (255, 0, 0), 2)\n",
    "    for num_txt in num_txt_2_person_list:\n",
    "        person_list = num_txt_2_person_list[num_txt]\n",
    "        t, l, b, r = person_list[0]\n",
    "        for person in person_list[1:]:\n",
    "            p_t, p_l, p_b, p_r = person\n",
    "            t = min(t, p_t)\n",
    "            l = min(l, p_l)\n",
    "            b = max(b, p_b)\n",
    "            r = max(r, p_r)\n",
    "        w = r-l\n",
    "        if(w>256):\n",
    "            print(\"ignore too wide person box\", l, \",\",r)\n",
    "            continue\n",
    "        num_person_list.append((num_txt, l, t, r, b))\n",
    "        cv2.putText(frame, num_txt, (l, t), cv2.FONT_HERSHEY_COMPLEX, 3, (255, 0, 0), 3)\n",
    "        cv2.rectangle(frame, (l, t), (r, b), (255, 0, 0), 2)\n",
    "    out_frame_id_2_person[frame_id] = num_person_list\n",
    "    writer.write(frame)\n",
    "    #im_show(frame)\n",
    "    \n",
    "reader.release()\n",
    "writer.release()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f7a463a",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_json = './fv_out.json'\n",
    "write_json_out(output_json, out_frame_id_2_person)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9313f52b",
   "metadata": {},
   "outputs": [],
   "source": [
    "out_frame_id_2_person[4500]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "335f90ee",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
